def testSingleTableQuery(self):
"""
Test a query on a single table (using chunk iterator)
"""
dbobj = DBObject(driver=self.driver, database=self.db_name)
query = 'SELECT id, sqrt FROM doubleTable'
results = dbobj.get_chunk_iterator(query)
dtype = [('id', int),
('sqrt', float)]
i = 1
for chunk in results:
for row in chunk:
self.assertEqual(row[0], i)
self.assertAlmostEqual(row[1], np.sqrt(i))
self.assertEqual(dtype, row.dtype)
i += 1
self.assertEqual(i, 201)
def testSingleTableQuery(self):
"""
Test a query on a single table (using chunk iterator)
"""
dbobj = DBObject(driver=self.driver, database=self.db_name)
query = 'SELECT id, sqrt FROM doubleTable'
results = dbobj.get_chunk_iterator(query)
dtype = [('id', int),
('sqrt', float)]
i = 1
for chunk in results:
for row in chunk:
self.assertEqual(row[0], i)
self.assertAlmostEqual(row[1], np.sqrt(i))
self.assertEqual(dtype, row.dtype)
i += 1
self.assertEqual(i, 201)
def testPassingConnection(self):
"""
Repeat the test from testJoin, but with a DBObject whose connection was passed
directly from another DBObject, to make sure that passing a connection works
"""
dbobj_base = DBObject(driver=self.driver, database=self.db_name)
dbobj = DBObject(connection=dbobj_base.connection)
query = 'SELECT doubleTable.id, intTable.id, doubleTable.log, intTable.thrice '
query += 'FROM doubleTable, intTable WHERE doubleTable.id = intTable.id'
results = dbobj.get_chunk_iterator(query, chunk_size=10)
dtype = [
('id', int),
('id_1', int),
('log', float),
('thrice', int)]
i = 0
for chunk in results:
if i < 90:
self.assertEqual(len(chunk), 10)
for row in chunk:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
self.assertEqual(dtype, row.dtype)
i += 1
self.assertEqual(i, 99)
# make sure that we found all the matches whe should have
results = dbobj.execute_arbitrary(query)
self.assertEqual(dtype, results.dtype)
i = 0
for row in results:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
i += 1
self.assertEqual(i, 99)
def testPassingConnection(self):
"""
Repeat the test from testJoin, but with a DBObject whose connection was passed
directly from another DBObject, to make sure that passing a connection works
"""
dbobj_base = DBObject(driver=self.driver, database=self.db_name)
dbobj = DBObject(connection=dbobj_base.connection)
query = 'SELECT doubleTable.id, intTable.id, doubleTable.log, intTable.thrice '
query += 'FROM doubleTable, intTable WHERE doubleTable.id = intTable.id'
results = dbobj.get_chunk_iterator(query, chunk_size=10)
dtype = [
('id', int),
('id_1', int),
('log', float),
('thrice', int)]
i = 0
for chunk in results:
if i < 90:
self.assertEqual(len(chunk), 10)
for row in chunk:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
self.assertEqual(dtype, row.dtype)
i += 1
self.assertEqual(i, 99)
# make sure that we found all the matches whe should have
results = dbobj.execute_arbitrary(query)
self.assertEqual(dtype, results.dtype)
i = 0
for row in results:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
i += 1
self.assertEqual(i, 99)
def testDtype(self):
"""
Test that passing dtype to a query works
(also test q query on a single table using .execute_arbitrary() directly
"""
dbobj = DBObject(driver=self.driver, database=self.db_name)
query = 'SELECT id, log FROM doubleTable'
dtype = [('id', int), ('log', float)]
results = dbobj.execute_arbitrary(query, dtype = dtype)
self.assertEqual(results.dtype, dtype)
for xx in results:
self.assertAlmostEqual(np.log(xx[0]), xx[1], 6)
self.assertEqual(len(results), 200)
results = dbobj.get_chunk_iterator(query, chunk_size=10, dtype=dtype)
next(results)
for chunk in results:
self.assertEqual(chunk.dtype, dtype)
def testDtype(self):
"""
Test that passing dtype to a query works
(also test q query on a single table using .execute_arbitrary() directly
"""
dbobj = DBObject(driver=self.driver, database=self.db_name)
query = 'SELECT id, log FROM doubleTable'
dtype = [('id', int), ('log', float)]
results = dbobj.execute_arbitrary(query, dtype = dtype)
self.assertEqual(results.dtype, dtype)
for xx in results:
self.assertAlmostEqual(np.log(xx[0]), xx[1], 6)
self.assertEqual(len(results), 200)
results = dbobj.get_chunk_iterator(query, chunk_size=10, dtype=dtype)
next(results)
for chunk in results:
self.assertEqual(chunk.dtype, dtype)
def testJoin(self):
"""
Test a join
"""
dbobj = DBObject(driver=self.driver, database=self.db_name)
query = 'SELECT doubleTable.id, intTable.id, doubleTable.log, intTable.thrice '
query += 'FROM doubleTable, intTable WHERE doubleTable.id = intTable.id'
results = dbobj.get_chunk_iterator(query, chunk_size=10)
dtype = [
('id', int),
('id_1', int),
('log', float),
('thrice', int)]
i = 0
for chunk in results:
if i < 90:
self.assertEqual(len(chunk), 10)
for row in chunk:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
self.assertEqual(dtype, row.dtype)
i += 1
self.assertEqual(i, 99)
# make sure that we found all the matches whe should have
results = dbobj.execute_arbitrary(query)
self.assertEqual(dtype, results.dtype)
i = 0
for row in results:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
i += 1
self.assertEqual(i, 99)
def testJoin(self):
"""
Test a join
"""
dbobj = DBObject(driver=self.driver, database=self.db_name)
query = 'SELECT doubleTable.id, intTable.id, doubleTable.log, intTable.thrice '
query += 'FROM doubleTable, intTable WHERE doubleTable.id = intTable.id'
results = dbobj.get_chunk_iterator(query, chunk_size=10)
dtype = [
('id', int),
('id_1', int),
('log', float),
('thrice', int)]
i = 0
for chunk in results:
if i < 90:
self.assertEqual(len(chunk), 10)
for row in chunk:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
self.assertEqual(dtype, row.dtype)
i += 1
self.assertEqual(i, 99)
# make sure that we found all the matches whe should have
results = dbobj.execute_arbitrary(query)
self.assertEqual(dtype, results.dtype)
i = 0
for row in results:
self.assertEqual(2*(i+1), row[0])
self.assertEqual(row[0], row[1])
self.assertAlmostEqual(np.log(row[0]), row[2], 6)
self.assertEqual(3*row[0], row[3])
i += 1
self.assertEqual(i, 99)
args = parser.parse_args()
if args.table is None or args.outdir is None:
raise RuntimeError("must specify table and out dir")
db = DBObject(database='LSSTCATSIM',
host='fatboy.phys.washington.edu',
port=1433,
driver='mssql+pymssql')
d_z = 50.0
out_dict = {}
dtype = np.dtype([('gal_l', float), ('gal_b', float), ('px', float)])
query = 'SELECT gal_l, gal_b, parallax FROM %s' % args.table
chunk_iter = db.get_chunk_iterator(query, chunk_size=100000, dtype=dtype)
for chunk in chunk_iter:
xyz = xyz_from_lon_lat_px(chunk['gal_l'], chunk['gal_b'],
0.001 * chunk['px'])
z_quant = np.round(np.abs(xyz[2] / d_z)).astype(int)
z_dex, z_ct = np.unique(z_quant, return_counts=True)
for dd, cc in zip(z_dex, z_ct):
if dd in out_dict:
out_dict[dd] += cc
else:
out_dict[dd] = cc
with open(os.path.join(args.outdir, "%s_z_dist.txt" % args.table),
'w') as output_file:
key_list = list(out_dict.keys())
